"""
增强数据下载器

支持批量下载、断点续传、多线程下载和进度监控
"""

import asyncio
import aiohttp
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional, Callable, Tuple
import pandas as pd
import time
import json
from pathlib import Path
import logging
from dataclasses import dataclass

from .data_manager import DataManager
from ..core.exceptions import NetworkError, DataMissingError
from ..utils.helpers import TimeUtils


@dataclass
class DownloadTask:
    """下载任务"""
    exchange: str
    symbol: str
    interval: str
    trade_type: str = 'spot'
    start_date: Optional[datetime] = None
    end_date: Optional[datetime] = None
    priority: int = 1  # 优先级，数字越小优先级越高
    
    def __post_init__(self):
        if not self.end_date:
            self.end_date = datetime.now()
        if not self.start_date:
            self.start_date = self.end_date - timedelta(days=30)


@dataclass
class DownloadProgress:
    """下载进度"""
    task_id: str
    exchange: str
    symbol: str
    interval: str
    total_chunks: int
    completed_chunks: int
    failed_chunks: int
    start_time: datetime
    estimated_time: Optional[float] = None
    
    @property
    def progress_percent(self) -> float:
        if self.total_chunks == 0:
            return 0.0
        return (self.completed_chunks / self.total_chunks) * 100
    
    @property
    def is_completed(self) -> bool:
        return self.completed_chunks == self.total_chunks
    
    @property
    def has_failed(self) -> bool:
        return self.failed_chunks > 0


class EnhancedDataDownloader:
    """增强数据下载器"""
    
    def __init__(self, data_manager: DataManager, config: Dict[str, Any] = None):
        self.data_manager = data_manager
        self.config = config or {}
        self.logger = logging.getLogger(__name__)
        
        # 下载配置
        self.max_workers = self.config.get('max_workers', 4)
        self.chunk_size_days = self.config.get('chunk_size_days', 7)
        self.retry_attempts = self.config.get('retry_attempts', 3)
        self.retry_delay = self.config.get('retry_delay', 1.0)
        self.rate_limit_delay = self.config.get('rate_limit_delay', 0.1)
        
        # 状态管理
        self.download_queue: List[DownloadTask] = []
        self.progress_tracker: Dict[str, DownloadProgress] = {}
        self.is_downloading = False
        self.download_stats = {
            'total_tasks': 0,
            'completed_tasks': 0,
            'failed_tasks': 0,
            'total_data_points': 0,
            'start_time': None,
            'end_time': None
        }
        
        # 回调函数
        self.progress_callback: Optional[Callable] = None
        self.completion_callback: Optional[Callable] = None
        
    def add_download_task(self, task: DownloadTask) -> str:
        """添加下载任务"""
        task_id = f"{task.exchange}_{task.symbol}_{task.interval}_{task.trade_type}_{int(time.time())}"
        task.task_id = task_id
        
        # 按优先级插入队列
        inserted = False
        for i, existing_task in enumerate(self.download_queue):
            if task.priority < existing_task.priority:
                self.download_queue.insert(i, task)
                inserted = True
                break
        
        if not inserted:
            self.download_queue.append(task)
        
        self.logger.info(f"添加下载任务: {task_id}")
        return task_id
    
    def add_batch_tasks(self, tasks: List[DownloadTask]) -> List[str]:
        """批量添加下载任务"""
        task_ids = []
        for task in tasks:
            task_id = self.add_download_task(task)
            task_ids.append(task_id)
        
        self.logger.info(f"批量添加 {len(tasks)} 个下载任务")
        return task_ids
    
    def create_symbol_tasks(self, exchange: str, symbols: List[str], 
                           intervals: List[str], trade_type: str = 'spot',
                           start_date: datetime = None, end_date: datetime = None) -> List[DownloadTask]:
        """为多个交易对创建下载任务"""
        tasks = []
        for symbol in symbols:
            for interval in intervals:
                task = DownloadTask(
                    exchange=exchange,
                    symbol=symbol,
                    interval=interval,
                    trade_type=trade_type,
                    start_date=start_date,
                    end_date=end_date
                )
                tasks.append(task)
        
        return tasks
    
    def start_download(self, progress_callback: Callable = None, 
                      completion_callback: Callable = None) -> bool:
        """开始下载"""
        if self.is_downloading:
            self.logger.warning("下载已在进行中")
            return False
        
        if not self.download_queue:
            self.logger.warning("下载队列为空")
            return False
        
        self.progress_callback = progress_callback
        self.completion_callback = completion_callback
        self.is_downloading = True
        
        # 重置统计信息
        self.download_stats = {
            'total_tasks': len(self.download_queue),
            'completed_tasks': 0,
            'failed_tasks': 0,
            'total_data_points': 0,
            'start_time': datetime.now(),
            'end_time': None
        }
        
        # 启动下载线程
        download_thread = threading.Thread(target=self._download_worker)
        download_thread.daemon = True
        download_thread.start()
        
        self.logger.info(f"开始下载 {len(self.download_queue)} 个任务")
        return True
    
    def _download_worker(self):
        """下载工作线程"""
        try:
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                # 提交所有任务
                future_to_task = {}
                for task in self.download_queue:
                    future = executor.submit(self._download_single_task, task)
                    future_to_task[future] = task
                
                # 处理完成的任务
                for future in as_completed(future_to_task):
                    task = future_to_task[future]
                    try:
                        success = future.result()
                        if success:
                            self.download_stats['completed_tasks'] += 1
                        else:
                            self.download_stats['failed_tasks'] += 1
                        
                        # 调用进度回调
                        if self.progress_callback:
                            self.progress_callback(self.get_overall_progress())
                            
                    except Exception as e:
                        self.logger.error(f"任务执行失败 {task.exchange}_{task.symbol}: {e}")
                        self.download_stats['failed_tasks'] += 1
        
        finally:
            self.is_downloading = False
            self.download_stats['end_time'] = datetime.now()
            
            # 调用完成回调
            if self.completion_callback:
                self.completion_callback(self.download_stats)
            
            self.logger.info("所有下载任务完成")
    
    def _download_single_task(self, task: DownloadTask) -> bool:
        """下载单个任务"""
        try:
            task_id = getattr(task, 'task_id', f"{task.exchange}_{task.symbol}_{task.interval}")
            
            # 计算时间块
            chunks = self._calculate_time_chunks(task.start_date, task.end_date)
            
            # 创建进度跟踪
            progress = DownloadProgress(
                task_id=task_id,
                exchange=task.exchange,
                symbol=task.symbol,
                interval=task.interval,
                total_chunks=len(chunks),
                completed_chunks=0,
                failed_chunks=0,
                start_time=datetime.now()
            )
            self.progress_tracker[task_id] = progress
            
            # 下载每个时间块
            all_data = []
            for i, (start_time, end_time) in enumerate(chunks):
                success = False
                for attempt in range(self.retry_attempts):
                    try:
                        # 检查是否已存在数据
                        existing_data = self._check_existing_data(
                            task.exchange, task.symbol, task.interval, 
                            task.trade_type, start_time, end_time
                        )
                        
                        if existing_data is not None and not existing_data.empty:
                            all_data.append(existing_data)
                            success = True
                            break
                        
                        # 下载新数据
                        chunk_data = self.data_manager.download_and_save_data(
                            exchange=task.exchange,
                            symbol=task.symbol,
                            interval=task.interval,
                            trade_type=task.trade_type,
                            start_time=start_time,
                            end_time=end_time
                        )
                        
                        if chunk_data:
                            success = True
                            break
                        
                    except Exception as e:
                        self.logger.warning(f"下载块失败 (尝试 {attempt + 1}): {e}")
                        if attempt < self.retry_attempts - 1:
                            time.sleep(self.retry_delay * (attempt + 1))
                
                if success:
                    progress.completed_chunks += 1
                else:
                    progress.failed_chunks += 1
                    self.logger.error(f"下载块最终失败: {task.symbol} {start_time}-{end_time}")
                
                # 速率限制
                time.sleep(self.rate_limit_delay)
            
            # 更新统计信息
            if all_data:
                total_data = pd.concat(all_data, ignore_index=True)
                self.download_stats['total_data_points'] += len(total_data)
            
            success_rate = progress.completed_chunks / progress.total_chunks
            self.logger.info(f"任务完成: {task_id}, 成功率: {success_rate:.1%}")
            
            return success_rate > 0.8  # 80%成功率认为任务成功
            
        except Exception as e:
            self.logger.error(f"下载任务失败 {task.exchange}_{task.symbol}: {e}")
            return False
    
    def _calculate_time_chunks(self, start_date: datetime, end_date: datetime) -> List[Tuple[datetime, datetime]]:
        """计算时间块"""
        chunks = []
        current_start = start_date
        chunk_delta = timedelta(days=self.chunk_size_days)
        
        while current_start < end_date:
            current_end = min(current_start + chunk_delta, end_date)
            chunks.append((current_start, current_end))
            current_start = current_end
        
        return chunks
    
    def _check_existing_data(self, exchange: str, symbol: str, interval: str,
                           trade_type: str, start_time: datetime, end_time: datetime) -> Optional[pd.DataFrame]:
        """检查已存在的数据"""
        try:
            # 尝试从数据管理器获取数据
            data = self.data_manager.get_data(
                exchange=exchange,
                symbol=symbol,
                interval=interval,
                trade_type=trade_type,
                start_date=start_time,
                end_date=end_time
            )
            
            if data is not None and not data.empty:
                # 检查数据完整性
                expected_points = self._calculate_expected_data_points(start_time, end_time, interval)
                actual_points = len(data)
                
                # 如果数据完整性超过90%，认为数据存在
                if actual_points >= expected_points * 0.9:
                    return data
            
            return None
            
        except Exception as e:
            self.logger.debug(f"检查已存在数据失败: {e}")
            return None
    
    def _calculate_expected_data_points(self, start_time: datetime, end_time: datetime, interval: str) -> int:
        """计算预期数据点数"""
        interval_minutes = {
            '1m': 1, '3m': 3, '5m': 5, '15m': 15, '30m': 30,
            '1h': 60, '2h': 120, '4h': 240, '6h': 360, '8h': 480, '12h': 720,
            '1d': 1440, '3d': 4320, '1w': 10080, '1M': 43200
        }
        
        minutes = interval_minutes.get(interval, 60)
        total_minutes = (end_time - start_time).total_seconds() / 60
        return int(total_minutes / minutes)
    
    def get_progress(self, task_id: str) -> Optional[DownloadProgress]:
        """获取任务进度"""
        return self.progress_tracker.get(task_id)
    
    def get_overall_progress(self) -> Dict[str, Any]:
        """获取整体进度"""
        total_tasks = self.download_stats['total_tasks']
        completed_tasks = self.download_stats['completed_tasks']
        failed_tasks = self.download_stats['failed_tasks']
        
        progress_percent = 0.0
        if total_tasks > 0:
            progress_percent = ((completed_tasks + failed_tasks) / total_tasks) * 100
        
        elapsed_time = None
        if self.download_stats['start_time']:
            end_time = self.download_stats['end_time'] or datetime.now()
            elapsed_time = (end_time - self.download_stats['start_time']).total_seconds()
        
        return {
            'total_tasks': total_tasks,
            'completed_tasks': completed_tasks,
            'failed_tasks': failed_tasks,
            'progress_percent': progress_percent,
            'is_downloading': self.is_downloading,
            'elapsed_time': elapsed_time,
            'total_data_points': self.download_stats['total_data_points']
        }
    
    def stop_download(self) -> bool:
        """停止下载"""
        if not self.is_downloading:
            return False
        
        self.is_downloading = False
        self.logger.info("下载已停止")
        return True
    
    def clear_queue(self) -> None:
        """清空下载队列"""
        self.download_queue.clear()
        self.progress_tracker.clear()
        self.logger.info("下载队列已清空")
    
    def get_queue_status(self) -> Dict[str, Any]:
        """获取队列状态"""
        return {
            'queue_size': len(self.download_queue),
            'is_downloading': self.is_downloading,
            'active_tasks': len([p for p in self.progress_tracker.values() if not p.is_completed]),
            'completed_tasks': len([p for p in self.progress_tracker.values() if p.is_completed]),
            'failed_tasks': len([p for p in self.progress_tracker.values() if p.has_failed])
        }
    
    def export_progress_report(self, output_path: str = None) -> str:
        """导出进度报告"""
        if not output_path:
            output_path = f"download_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        report = {
            'download_stats': self.download_stats,
            'overall_progress': self.get_overall_progress(),
            'task_progress': {
                task_id: {
                    'exchange': progress.exchange,
                    'symbol': progress.symbol,
                    'interval': progress.interval,
                    'progress_percent': progress.progress_percent,
                    'completed_chunks': progress.completed_chunks,
                    'total_chunks': progress.total_chunks,
                    'failed_chunks': progress.failed_chunks,
                    'is_completed': progress.is_completed,
                    'has_failed': progress.has_failed
                }
                for task_id, progress in self.progress_tracker.items()
            },
            'queue_status': self.get_queue_status(),
            'config': self.config,
            'generated_at': datetime.now().isoformat()
        }
        
        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(report, f, indent=2, ensure_ascii=False)
            
            self.logger.info(f"进度报告已导出: {output_path}")
            return output_path
            
        except Exception as e:
            self.logger.error(f"导出进度报告失败: {e}")
            return ""


def create_common_download_tasks(exchanges: List[str] = None, 
                               symbols: List[str] = None,
                               intervals: List[str] = None,
                               days_back: int = 30) -> List[DownloadTask]:
    """创建常用的下载任务"""
    if not exchanges:
        exchanges = ['okx', 'binance']
    
    if not symbols:
        symbols = ['BTC-USDT', 'ETH-USDT', 'BNB-USDT', 'ADA-USDT', 'SOL-USDT']
    
    if not intervals:
        intervals = ['1m', '5m', '15m', '1h', '4h', '1d']
    
    end_date = datetime.now()
    start_date = end_date - timedelta(days=days_back)
    
    tasks = []
    for exchange in exchanges:
        for symbol in symbols:
            for interval in intervals:
                # 现货任务
                task_spot = DownloadTask(
                    exchange=exchange,
                    symbol=symbol,
                    interval=interval,
                    trade_type='spot',
                    start_date=start_date,
                    end_date=end_date,
                    priority=1 if interval in ['1h', '4h', '1d'] else 2
                )
                tasks.append(task_spot)
                
                # 合约任务
                task_futures = DownloadTask(
                    exchange=exchange,
                    symbol=symbol,
                    interval=interval,
                    trade_type='futures',
                    start_date=start_date,
                    end_date=end_date,
                    priority=1 if interval in ['1h', '4h', '1d'] else 2
                )
                tasks.append(task_futures)
    
    return tasks